#include "cytoskel-solver.hpp"
#include "SimTKsimbody_aux.h"

using namespace SimTK;

CytoskelSubsystemImpl::CytoskelSubsystemImpl()
  : ForceSubsystem::Guts("Cytoskel", "1.0") 
{}

Subsystem::Guts* 
CytoskelSubsystemImpl::cloneImpl() const
{ return new CytoskelSubsystemImpl(); }

int
CytoskelSubsystemImpl::realizeSubsystemTopologyImpl(State& state) const
{
  Array_<Connection>* c = new Array_<Connection>();
  connectionsIdx = allocateDiscreteVariable(state, Stage::Dynamics,
					    new Value<Array_<Connection>*>(c));
}

int
CytoskelSubsystemImpl::realizeSubsystemDynamicsImpl(const State& state) const 
{
  const MultibodySystem& system = MultibodySystem::downcast(getSystem());
  const SimbodyMatterSubsystem& matter = system.getMatterSubsystem();
  Vector_<SpatialVec>& forces = system.updRigidBodyForces(state,
							  Stage::Dynamics);
  Array_<Connection>* connections = getConnections(state);
  for (Array_<Connection>::iterator p = connections->begin();
       p != connections->end();
       ++p)
    {
      MobilizedBodyIndex i = p->body1->getMobilizedBodyIndex();
      MobilizedBodyIndex j = p->body2->getMobilizedBodyIndex();
      /*const MobilizedBody& body1 = matter.getMobilizedBody(i);
	const MobilizedBody& body2 = matter.getMobilizedBody(j);*/
      
      Vec3 r = p->body1->getBodyOriginLocation(state)-
	p->body2->getBodyOriginLocation(state);
      Real distance = r.norm();

      Vec3 u = p->body1->getBodyOriginVelocity(state) - p->body2->getBodyOriginVelocity(state);

      Vec3 force = -0.5 * ((distance - p->l) / distance) * r - 0.25 * u;
      forces[i][1] += force;
      forces[j][1] -= force;
    }
  return 0;
}

Real
CytoskelSubsystemImpl::calcPotentialEnergy(const State& state) const 
{}

void
CytoskelSubsystemImpl::setConnections(State& state, Array_<Connection>* c)
{
  Value<Array_<Connection>*>::updDowncast(updDiscreteVariable(state,
							      connectionsIdx)) = c;
}

Array_<Connection>*
CytoskelSubsystemImpl::getConnections(const State& state) const
{
  return Value<Array_<Connection>*>::downcast(getDiscreteVariable(state,
								  connectionsIdx));
}

CytoskelSubsystem::CytoskelSubsystem (MultibodySystem& system)
{ 
  adoptSubsystemGuts(new CytoskelSubsystemImpl());
  system.addForceSubsystem(*this);
}

#if 1

int
main (int argv, char** argc)
{
  MultibodySystem system;
  SimbodyMatterSubsystem matter(system);
  CytoskelSubsystem cytoskel(system);
  Body::Rigid body (MassProperties(1.0, Vec3(0), Inertia(1)));
  body.addDecoration(Transform(), DecorativeSphere(0.1));
  MobilizedBody::Free body1 (matter.Ground(), body);
  MobilizedBody::Free body2 (matter.Ground(), body);
  MobilizedBody::Free body3 (matter.Ground(), body);
  system.updDefaultSubsystem().addEventReporter(new VTKEventReporter(system, 0.01));
  
  body2.setDefaultTranslation (Vec3(1.5, 0, 0));
  body3.setDefaultTranslation (Vec3(0, 1.5, 0));
  
  system.realizeTopology();
  State state = system.getDefaultState();
  Array_<Connection>* cs = cytoskel.getConnections(state);
  cs->push_back ({1, 2, &body1, &body2, 1});
  cs->push_back ({1, 3, &body1, &body3, 1});
  cs->push_back ({2, 3, &body2, &body3, 1});
  cytoskel.setConnections(state, cs);

  RungeKuttaMersonIntegrator integ(system);
  TimeStepper ts(system, integ);
  ts.initialize(state);
  ts.stepTo(50.0);

  return 0;
}

#endif
